import numpy as np
from load_data import read_data, gen_dict


def crf_viterbi(tk, sl, tag_num, n):
    #条件随机场的viterbi算法
    '''
    :param tk:定义在边上的特征函数，ndarray,
    :param sl: 定义在节点上的特征函数，ndarray,
    :param l: 随机变量可以取的值，ndarray,
    :param n: 观测序列的长度,int,也就是位置的个数
    :return: 最优的标记序列，由l里面的值组成的数组,ndarray
    '''

    #第一步：初始化：
    #定义最优标记序列
    best_tag_list = []
    delta = np.zeros([n, tag_num])
    for i in range(tag_num):
        delta[0][i] = sl[0][i]
    best_tag_list.append(np.argmax(delta[0]))

    #第二步递推：
    for i in range(1, n): #遍历位置，当前位置i，上一位置i-1
        for l in range(tag_num): #遍历可选的标签
            #计算的是delta[i][l],因为取最大非规范化概率，需要遍历
            max = 0
            for j in range(tag_num):
                #使用tk中的第 i - 1个矩阵概率
                #使用sl中的当前第 i 个结点的概率
                p = delta[i-1][j] + tk[i-1][j][l] + sl[i][l];
                if p > max:
                    max = p
            delta[i][l] = max
        best_tag_list.append(np.argmax(delta[i]))
    return np.array(best_tag_list)

train_sentence_word_ids_list = None
train_sentence_tags_list = None

def crf_train_tk(train_data, id1, id2, tag_num):
    global train_sentence_word_ids_list
    global train_sentence_tags_list

    tk = np.zeros([tag_num, tag_num])
    if train_sentence_word_ids_list is None:
        train_sentence_word_ids_list, train_sentence_tags_list = read_data(train_data, tag_to_ids, dict)
    sentence_num = len(train_sentence_word_ids_list)
    for i in range(sentence_num):
        sentence_length = len(train_sentence_word_ids_list[i])
        if sentence_length <= 1:
            continue
        for j in range(sentence_length - 1):
            if train_sentence_word_ids_list[i][j] == id1 and train_sentence_word_ids_list[i][j+1] == id2:
                tk[train_sentence_tags_list[i][j]][train_sentence_tags_list[i][j+1]] += 1
    return tk

def crf_train_sl(train_data, tag_to_ids, dict):
    sentence_word_ids_list, sentence_tags_list = read_data(train_data, tag_to_ids, dict)
    tag_num = len(tag_to_ids)
    d = len(dict)
    sl = np.zeros([d, tag_num])
    sentence_num = len(sentence_word_ids_list)
    for i in range(sentence_num):
        sentence_length = len(sentence_word_ids_list[i])
        if sentence_length <= 1:
            continue
        for j in range(sentence_length - 1):
            sl[sentence_word_ids_list[i][j]][sentence_tags_list[i][j]] += 1
        sl[sentence_word_ids_list[i][sentence_length - 1]][sentence_tags_list[i][sentence_length - 1]] += 1
    return sl

def crf_test(sl, train_data, test_data, tag_to_ids, dict):

    sentence_word_ids_list, sentence_tags_list = read_data(test_data, tag_to_ids, dict)
    tag_num = len(tag_to_ids)


    sentence_num = len(sentence_word_ids_list)
    acc = 0
    total = 0
    fn = 0
    for i in range(sentence_num):

        sentence_length = len(sentence_word_ids_list[i])
        if sentence_length <= 1:
            continue
        test_sl = np.zeros([sentence_length, tag_num])
        # 这个句子有sentence_length个单词，只有 sentence_length - 1个 tag_num * tag_num的转移矩阵
        test_tk = np.zeros([sentence_length - 1, tag_num, tag_num])
        for j in range(sentence_length):
            test_sl[j] = sl[sentence_word_ids_list[i][j]]

        for j in range(1, sentence_length):
            tk = crf_train_tk(train_data, sentence_word_ids_list[i][j-1], sentence_word_ids_list[i][j], tag_num)
            test_tk[j-1] = tk

        best_tag_list = crf_viterbi(test_tk, test_sl, tag_num, sentence_length)

        total += len(sentence_tags_list[i])
        acc += np.sum(best_tag_list == np.array(sentence_tags_list[i]))
        if i % 50 == 0:
            print("进度: {}%".format(i / sentence_num * 100))
    print("正确率为:{}".format(acc / total))

if __name__ == "__main__":
    train_data = "./data/train"
    testa_data = "./data/testa"
    testb_data = "./data/testb"
    tag_to_ids = {"B-ORG": 0, "I-ORG": 1, "B-PER": 2, "I-PER": 3, "B-LOC": 4, "I-LOC": 5, "B-MISC": 6, "I-MISC": 7,
                  "O": 8}

    dict = gen_dict([train_data, testa_data, testb_data])
    sl = crf_train_sl([train_data], tag_to_ids, dict)
    crf_test(sl, [train_data], [testa_data], tag_to_ids, dict)